package warden

import (
	"context"
	"fmt"

	"github.com/dghubble/trie"

	"google.golang.org/grpc"
)

type Router struct {
	t *trie.PathTrie
}

func NewRouter() *Router {
	return &Router{t: trie.NewPathTrie()}
}

func (r *Router) Add(method string, handlers ...grpc.UnaryServerInterceptor) {
	if len(handlers) == 0 {
		panic(fmt.Sprintf("warden: adding routes(%s) handlers is empty!", method))
	}
	if r.t.Get(method) != nil {
		panic(fmt.Sprintf("warden: Repeat adding routes(%s)!", method))
	}
	r.t.Put(method, setupInterceptor(handlers...))
}

func (r *Router) Get(method string) grpc.UnaryServerInterceptor {
	value := r.t.Get(method)
	if value != nil {
		return value.(grpc.UnaryServerInterceptor)
	}
	return nil
}

func (r *Router) Handle(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
	interceptor := r.Get(args.FullMethod)
	if interceptor != nil {
		return interceptor(ctx, req, args, handler)
	}
	return handler(ctx, req)
}

func setupInterceptor(handlers ...grpc.UnaryServerInterceptor) (inter grpc.UnaryServerInterceptor) {
	n := len(handlers)
	if n == 0 {
		return
	} else if n == 1 {
		inter = handlers[0]
		return
	}
	inter = func(ctx context.Context, req interface{}, args *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		var (
			i     int
			chain grpc.UnaryHandler
		)

		chain = func(ic context.Context, ir interface{}) (interface{}, error) {
			if i == n-1 {
				return handler(ic, ir)
			}
			i++
			return handlers[i](ic, ir, args, chain)
		}

		return handlers[0](ctx, req, args, chain)
	}
	return
}
